#####################################################################
# Figure 8
#####################################################################

#####################################################################
# Preliminaries
#####################################################################

rm(list = ls())

library(limSolve)
library(scinference)

set.seed(12345)

###################################################################
# Functions
###################################################################

source("code/auxiliary scripts/common_functions.R")

sim <- function(T0,T1,J,K,rho_u,w0_sc,var_param,dgp){
  
  # Preliminaries
  sig_level <- 0.1
  w0_did    <- rep(1/J,J)
  
  # Generate T01
  T01 <- T0+T1
  
  # DGPs
  
  Y0  <- matrix(NA,T01,J)
  for (j in 1:J){
    Y0[,j] <- rnorm(T01,mean=j,sd=1)
  }
  
  u.sim <- generate.AR.series(T01,rho_u)
  
  # DGP1: Markov-switching weights
  if (dgp==1){
    if (var_param == 0){
      w  <- w0_sc
      Y1 <- Y0 %*% w + u.sim 
    }
    if (var_param>0){
      Y1 <- matrix(NA,T01,1)
      s_realized  <- 1
      for (t in 1:T01){
        regime_stay <- (rbinom(1,1,var_param)==1)
        s_realized  <- ifelse(regime_stay,s_realized,1-s_realized)
        w_realized  <- cbind(w0_sc*s_realized+w0_did*(1-s_realized))
        Y1[t]       <- rbind(Y0[t,]) %*% w_realized + u.sim[t]
      }
    }
  }
  
  # DGP2: Adding normal noise to the weights
  if (dgp==2){
    if (var_param == 0){
      w  <- w0_sc
      Y1 <- Y0 %*% w + u.sim
    }
    if (var_param>0){
      Y1 <- matrix(NA,T01,1)
      for (t in 1:T01){
        w_noise <- w0_sc + rnorm(J,0,sd=var_param)
        Y1[t]   <- rbind(Y0[t,]) %*% w_noise + u.sim[t]
      }
    }
  }
  
  # SC with debiasing
  r_sc_sim    <- scinference(Y1,Y0,T1=T1,T0=T0,inference_method="ttest",estimation_method="sc",K=K,lsei_type=1)
  tau_sc      <- r_sc_sim$att
  se_sc       <- r_sc_sim$se
  cov_sc      <- (abs(tau_sc/se_sc)<=qt(1-sig_level/2,df=K-1))
  leng_sc     <- 2*qt(1-sig_level/2,df=K-1)*se_sc
  
  return(c(cov_sc,leng_sc))
  
}

#####################################################################
# Calibration DGPs
#####################################################################

source("code/auxiliary scripts/calibration_dgps.R")

###################################################################
# Simulations
###################################################################

# Setup

nreps <- 50000

var_param_vec_dgp1 <- seq(0,0.9,0.1)
var_param_vec_dgp2 <- seq(0,0.5,0.05)

# Simulation loops

results_overall_dgp1 <- matrix(NA,length(var_param_vec_dgp1),2)
for (v in 1:length(var_param_vec_dgp1)){
  results_temp_dgp1 <- matrix(NA,nreps,2)
  for (r in 1:nreps){
    results_temp_dgp1[r,] <- sim(T0=T0_app,T1=T1_app,J=J_app,K=3,rho_u=rho_u,w0_sc=w0_sc,var_param=var_param_vec_dgp1[v],dgp=1)
  }
  results_overall_dgp1[v,] <- colMeans(results_temp_dgp1)
}

results_overall_dgp2 <- matrix(NA,length(var_param_vec_dgp2),2)
for (v in 1:length(var_param_vec_dgp2)){
  results_temp_dgp2 <- matrix(NA,nreps,2)
  for (r in 1:nreps){
    results_temp_dgp2[r,] <- sim(T0=T0_app,T1=T1_app,J=J_app,K=3,rho_u=rho_u,w0_sc=w0_sc,var_param=var_param_vec_dgp2[v],dgp=2)
  }
  results_overall_dgp2[v,] <- colMeans(results_temp_dgp2)
}


###################################################################
# Figures
###################################################################

# Preparing figures

coverage_dgp1 <- results_overall_dgp1[,1]
coverage_dgp2 <- results_overall_dgp2[,1]

rel_length_dgp1 <- results_overall_dgp1[,2]/c(results_overall_dgp1[1,2])*100
rel_length_dgp2 <- results_overall_dgp2[,2]/c(results_overall_dgp2[1,2])*100

# Coverage

graphics.off()
pdf("graphics/time_varying_weights_coverage_correct_specification.pdf",pointsize=16,width=8.0,height=6.0)
plot(NULL, xlab=expression(paste("Probability of staying in the same regime ",(p[stay]))), ylab="Coverage", main = "Correct specification: coverage",col="blue", pch=".", xlim = range(var_param_vec_dgp1[-1]), ylim=c(0.4,1), xaxt='n')
lines(var_param_vec_dgp1[-1],coverage_dgp1[-1],col="black",lwd=5, lty=1,type="b",pch=1)
abline(h=0.9,lty=1,lwd=1,col="gray")
axis(1, at = var_param_vec_dgp1[-1])
dev.off()

graphics.off()
pdf("graphics/time_varying_weights_coverage_misspecification.pdf",pointsize=16,width=8.0,height=6.0)
plot(NULL, xlab=expression(paste("Variability in the weights ",(sigma[w]))), main = "Misspecification: coverage", ylab="Coverage", col="blue", pch=".", xlim = range(var_param_vec_dgp2[-1]), ylim=c(0.4,1), xaxt='n')
lines(var_param_vec_dgp2[-1],coverage_dgp2[-1],col="black",lwd=5, lty=1,type="b",pch=1)
abline(h=0.9,lty=1,lwd=1,col="gray")
axis(1, at = var_param_vec_dgp2[-1])
dev.off()

# Length

graphics.off()
pdf("graphics/time_varying_weights_length_correct_specification.pdf",pointsize=16,width=8.0,height=6.0)
plot(NULL, xlab=expression(paste("Probability of staying in the same regime ",(p[stay]))), ylab="Relative length (in %)", main="Correct specification: relative length", col="blue", pch=".", xlim = range(var_param_vec_dgp1[-1]), ylim=range(rel_length_dgp1), xaxt='n')
lines(var_param_vec_dgp1[-1],rel_length_dgp1[-1],col="black",lwd=5, lty=1,type="b",pch=1)
axis(1, at = var_param_vec_dgp1[-1])
dev.off()

graphics.off()
pdf("graphics/time_varying_weights_length_misspecification.pdf",pointsize=16,width=8.0,height=6.0)
plot(NULL, xlab=expression(paste("Variability in the weights ",(sigma[w]))), main = "Misspecification: relative length", ylab="Relative length (in %)", col="blue", pch=".", xlim = range(var_param_vec_dgp2[-1]), ylim=range(rel_length_dgp2), xaxt='n')
lines(var_param_vec_dgp2[-1],rel_length_dgp2[-1],col="black",lwd=5, lty=1,type="b",pch=1)
axis(1, at = var_param_vec_dgp2[-1])
dev.off()
